from langgraph.graph import StateGraph,START,END
from typing import TypedDict

from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate

import getpass
import os


api_key = "sk-6S0PtpNia71gjcfwSsDPsJ9mGqsVPr2XRQzAx1dHbJS7RW4t"
api_base="https://chatapi.littlewheat.com/v1"

class InputState(TypedDict):
    question:str


class OutputState(TypedDict):
    answer:str

class OverallState(InputState,OutputState):
    pass

def llm_node(state:InputState):
    messages=[
        {"role":"system","content":"你是一位乐于助人的智能小助理"},
        {"role":"user","content":state["question"]}
    ]

    llm = ChatOpenAI(model="gpt-4o",api_key=api_key,base_url = api_base,temperature=0)
    response = llm.invoke(messages)
    # print(response)
    return {"question":response.content}

def action_node(state:InputState):
    print(state)
    messages = [
        {"role":"system","content":"无论你接收到什么语言的文本，请翻译成英语"},
        {"role":"user","content":state["question"]}
    ]

    llm = ChatOpenAI(model="gpt-4o",api_key=api_key,base_url = api_base,temperature=0)
    response = llm.invoke(messages)
    print(response)
    return {"answer":response.content}

builder = StateGraph(state_schema=OverallState,input=InputState,output=OutputState)

# 添加节点
builder.add_node("llm_node",llm_node)
builder.add_node("action_node",action_node)

# 添加边
builder.add_edge(START,"llm_node")
builder.add_edge("llm_node","action_node")
builder.add_edge("action_node",END)

graph = builder.compile()

res = graph.invoke({"question":"你好，请你详细的介绍一下你自己"})
print(res)
